{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation data size: 316 316 316\n",
      "Train data size: 800 800 400\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import cv2\n",
    "import time\n",
    "import random\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from keras.models import model_from_json\n",
    "import shutil\n",
    "from utils import load_img, gen_x_y, eval_loss, gen_paths, ssim_eucli_loss, random_cropping, local_sum_loss, fix_singular_shape\n",
    "\n",
    "\n",
    "# Settings\n",
    "dataset = \"B\"\n",
    "with_validation = False\n",
    "rand_crop = True\n",
    "%matplotlib inline\n",
    "\n",
    "img_paths_test, img_paths_train = gen_paths(\n",
    "    path_file_root='data/paths_train_val_test',\n",
    "    dataset=dataset,\n",
    "    with_validation=with_validation\n",
    ")\n",
    "img_paths_test = sorted(img_paths_test, key=lambda x: int(x.split('_')[-1].split('.')[0]))\n",
    "# img_paths_test, img_paths_train = img_paths_test[:10], img_paths_train[:10]\n",
    "if with_validation:\n",
    "    img_paths_train = list(set(img_paths_train) - set(img_paths_val))\n",
    "    x_val, y_val, img_paths_val = gen_x_y(img_paths_val, 'val')\n",
    "    print(len(x_val), len(y_val), len(img_paths_val))\n",
    "    x_test, y_test, img_paths_test = gen_x_y(img_paths_test, 'test')\n",
    "    print('Test data size:', len(x_test), len(y_test), len(img_paths_test))\n",
    "else:\n",
    "    x_val, y_val, img_paths_val = gen_x_y(img_paths_test[:], 'test')\n",
    "    print('Validation data size:', len(x_val), len(y_val), len(img_paths_val))\n",
    "x_train, y_train, img_paths_train = gen_x_y(img_paths_train[:], 'train', augmentation_methods=['ori', 'flip'])\n",
    "print('Train data size:', len(x_train), len(y_train), len(img_paths_train))\n",
    "weights_dir = 'weights_' + dataset\n",
    "if os.path.exists(weights_dir):\n",
    "    shutil.rmtree(weights_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from keras.optimizers import Adam\n",
    "from utils import ssim_eucli_loss\n",
    "from SANet import SANet\n",
    "\n",
    "\n",
    "# Settings\n",
    "dataset = 'B'\n",
    "weights_dir = 'weights_' + dataset\n",
    "if os.path.exists(weights_dir):\n",
    "    shutil.rmtree(weights_dir)\n",
    "os.makedirs(weights_dir)\n",
    "net = 'SANet'\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "LOSS = 'ssim_eucli_loss'\n",
    "lr = 1e-4\n",
    "input_shape = (None, None, 3)#np.array((768, 1024, 3), dtype=int).tolist()\n",
    "\n",
    "model = SANet(input_shape=input_shape, IN=False)\n",
    "model.summary()\n",
    "optimizer = Adam(lr=lr)\n",
    "model.compile(optimizer=optimizer, loss=ssim_eucli_loss)\n",
    "from keras.utils import plot_model\n",
    "plot_model(model, 'models/{}.png'.format(net))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.load_weights(os.path.join(weights_dir, 'SANet_best.hdf5'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Settings\n",
    "lossMAE = 1e5\n",
    "lossMDMD, lossMAPE, lossMSE = -1, -1, -1\n",
    "counter_train = 0\n",
    "mae = 1e5\n",
    "mdmd = 1e5\n",
    "val_rate = 0.25\n",
    "lossesMDMD = []\n",
    "lossesMAE = []\n",
    "lossesMAPE = []\n",
    "lossesMSE = []\n",
    "path_val_display = img_paths_val[0]\n",
    "x_val_display = load_img(path_val_display)\n",
    "y_val_display = np.squeeze(y_val[0])\n",
    "x_val_display, y_val_display = fix_singular_shape(x_val_display), fix_singular_shape(y_val_display)\n",
    "with open('./models/{}.json'.format(net), 'w') as fout:\n",
    "    fout.write(model.to_json())\n",
    "if dataset == 'A':\n",
    "    save_frequencies = [(90, 0.1), (80, 0.05), (95, 0.25)]\n",
    "else:\n",
    "    save_frequencies = [(20, 0.1), (15, 0.05), (22, 0.25)]\n",
    "# epoch_train = 100\n",
    "\n",
    "# Set differernt modules trainable 3 Inception-like module and the Decoder module trainable seperatly.\n",
    "branches_trainable = [\n",
    "    list(range(1, 7+1)),\n",
    "    list(range(8, 14+1)),\n",
    "    list(range(15, 21+1)),\n",
    "    list(range(22, 28+1)),\n",
    "    list(range(29, 34+1))\n",
    "]\n",
    "if_train_seperatly = True\n",
    "if if_train_seperatly:\n",
    "    module_step = 40\n",
    "    epoch_train_seperatly = module_step * len(branches_trainable) * 1\n",
    "    epoch_train = epoch_train_seperatly + 100\n",
    "else:\n",
    "    epoch_train = 200\n",
    "\n",
    "# Training\n",
    "time_st = time.time()\n",
    "for epoch in range(epoch_train):\n",
    "    if if_train_seperatly:\n",
    "        trainable_choice = int(epoch/module_step) # epoch % (5*len(branches_trainable))\n",
    "        print('trainable_choice:', trainable_choice)\n",
    "        # Train modules seperatly to avoid output all zeros\n",
    "        if epoch < epoch_train_seperatly and trainable_choice in list(range(len(branches_trainable))):# [0, module_step*1, module_step*2, module_step*3, module_step*4, module_step*5]:\n",
    "            branch_trainable = branches_trainable[trainable_choice]\n",
    "            for i in range(1, len(model.layers)):\n",
    "                if 'conv' in model.layers[i].name:\n",
    "                    model.layers[i].trainable = False\n",
    "            for i in range(1, len(model.layers)):\n",
    "                idx_operator = int(model.layers[i].name.split('_')[-1])\n",
    "                if 'conv' in model.layers[i].name and idx_operator in branch_trainable:\n",
    "                    model.layers[i].trainable = True\n",
    "            model.compile(optimizer=Adam(lr=lr), loss=ssim_eucli_loss)\n",
    "        elif epoch == epoch_train_seperatly:\n",
    "            for i in range(1, len(model.layers)):\n",
    "                model.layers[i].trainable = True\n",
    "            model.compile(optimizer=Adam(lr=lr/1), loss=ssim_eucli_loss)\n",
    "        else:\n",
    "            pass\n",
    "\n",
    "#     if epoch == 35:\n",
    "#         for i in range(1, 14):\n",
    "#         #     print(model.layers[i].name)\n",
    "#             model.layers[i].trainable = True\n",
    "#         optimizer = Adam(lr=lr/10)\n",
    "#         model.compile(optimizer=optimizer, loss=local_sum_loss)\n",
    "    for i in range(len(x_train)):\n",
    "        if lossMAE < save_frequencies[0][0]:\n",
    "            val_rate = save_frequencies[0][1]\n",
    "            if lossMAE < save_frequencies[1][0]:\n",
    "                val_rate = save_frequencies[1][1]\n",
    "        if len(lossesMAE) > 100 and val_rate <= save_frequencies[-1][-1] and np.median(lossesMAE[-20]) > save_frequencies[-1][0]:\n",
    "            val_rate = save_frequencies[-1][-1]\n",
    "        x_, y_ = x_train[i], y_train[i]\n",
    "        model.fit(x_, y_, verbose=0)\n",
    "#         x_cropped, y_cropped = random_cropping(x_, y_)\n",
    "#         for idx_fit in range(x_cropped.shape[0]):\n",
    "#             model.fit(x_cropped[idx_fit:idx_fit+1], y_cropped[idx_fit:idx_fit+1], verbose=0)\n",
    "        counter_train += 1\n",
    "#         if epoch < 5:\n",
    "#             val_rate = 0.1\n",
    "        if counter_train % int(len(x_train)*val_rate) == 0:\n",
    "            # Calc loss\n",
    "            lossMDMD, lossMAE, lossMAPE, lossMSE = eval_loss(model, x_val, y_val)\n",
    "            lossesMDMD.append(lossMDMD)\n",
    "            lossesMAE.append(lossMAE)\n",
    "            lossesMAPE.append(lossMAPE)\n",
    "            lossesMSE.append(lossMSE)\n",
    "            lossMAE, lossMAPE, lossMDMD, lossMSE = round(lossMAE, 3), round(lossMAPE, 3), round(lossMDMD, 3), round(lossMSE, 3)\n",
    "            if ((lossMAE < mae or lossMDMD < mdmd) and epoch >= 0) or lossMAE < save_frequencies[1][0] * 0.9:\n",
    "                model.save_weights(\n",
    "                    os.path.join(weights_dir, '{}_MAE{}_MSE{}_MAPE{}_MDMD{}_epoch{}-{}.hdf5'.format(\n",
    "                        net, str(lossMAE), str(lossMSE), str(lossMAPE), str(lossMDMD), epoch, (counter_train-1)%len(x_train)\n",
    "                    ))\n",
    "                )\n",
    "                if lossMAE < mae:\n",
    "                    mae = lossMAE\n",
    "                    model.save_weights(os.path.join(weights_dir, '{}_best.hdf5'.format(net)))\n",
    "                if lossMDMD < mdmd:\n",
    "                    mdmd = lossMDMD\n",
    "            if counter_train % (len(x_train)*1) == 0:\n",
    "                # show prediction\n",
    "                pred = np.squeeze(model.predict(np.expand_dims(x_val_display, axis=0)))\n",
    "                fg, (ax_x_ori, ax_y, ax_pred) = plt.subplots(1, 3, figsize=(20, 4))\n",
    "                ax_x_ori.imshow(cv2.cvtColor(cv2.imread(path_val_display), cv2.COLOR_BGR2RGB))\n",
    "                ax_x_ori.set_title('Original Image')\n",
    "                ax_y.imshow(y_val_display, cmap=plt.cm.jet)\n",
    "                ax_y.set_title('Ground_truth: ' + str(np.sum(y_val_display)))\n",
    "                ax_pred.imshow(pred, cmap=plt.cm.jet)\n",
    "                ax_pred.set_title('Prediction: ' + str(np.sum(pred)))\n",
    "                plt.suptitle('Loss = ' + str(lossMAE))\n",
    "                plt.show()\n",
    "                if counter_train % (len(x_train)*4) == 0 or (epoch+1) % (20) == 0:\n",
    "                    # plot val_loss\n",
    "                    plt.plot(lossesMAPE, 'y')\n",
    "                    plt.plot(lossesMAE, 'b')\n",
    "                    plt.plot(lossesMSE, 'c')\n",
    "                    plt.plot(lossesMDMD, 'r')\n",
    "                    plt.legend(['LossMAPE', 'LossMAE', 'LossMSE', 'Loss_Density_Map_Distance'])\n",
    "                    plt.title('Loss')\n",
    "                    plt.show()\n",
    "        time_now = time.time()\n",
    "        time_consuming = time_now - time_st\n",
    "        sys.stdout.write('In epoch {}_{}, with MAE={}, MSE={}, MAPE={}, MDMD={}, time consuming={}m-{}s\\r'.format(\n",
    "            epoch, counter_train%len(x_train), str(lossMAE), str(lossMSE), str(lossMAPE), str(lossMDMD),\n",
    "            int(time_consuming/60), int(time_consuming-int(time_consuming/60)*60)\n",
    "        ))\n",
    "        sys.stdout.flush()\n",
    "end_time_of_train = '-'.join(time.ctime().split()[:-2])\n",
    "suffix_new_dir = '_{}_{}_bestMAE{}_{}'.format(dataset, LOSS, str(round(mae, 3)), end_time_of_train)\n",
    "shutil.move('weights_{}'.format(dataset), 'weights'+suffix_new_dir)\n",
    "\n",
    "loss_dir = 'losses_' + dataset\n",
    "if not os.path.exists(loss_dir):\n",
    "    os.makedirs(loss_dir)\n",
    "np.savetxt(os.path.join(loss_dir, 'loss_DMD.txt'), lossesMDMD)\n",
    "np.savetxt(os.path.join(loss_dir, 'loss_MAE.txt'), lossesMAE)\n",
    "np.savetxt(os.path.join(loss_dir, 'loss_MAPE.txt'), lossesMAPE)\n",
    "np.savetxt(os.path.join(loss_dir, 'loss_MSE.txt'), lossesMSE)\n",
    "plt.plot(lossesMAPE, 'y')\n",
    "plt.plot(lossesMAE, 'b')\n",
    "plt.plot(lossesMSE, 'c')\n",
    "plt.plot(lossesMDMD, 'r')\n",
    "plt.legend(['LossMAPE', 'LossMAE', 'LossMSE', 'Loss_Density_Map_Distance'])\n",
    "plt.title('Loss -- {} epochs'.format(epoch_train))\n",
    "plt.savefig('./losses_{}/loss_{}_{}.jpg'.format(dataset, dataset, end_time_of_train))\n",
    "plt.show()\n",
    "shutil.move('losses_{}'.format(dataset), 'losses'+suffix_new_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
